from vanna.flask import VannaFlaskApp
from langchain_zhipu import ChatZhipuAI
from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from vanna_auth import SimplePassword

client = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                     , model="glm-4")


class MyVanna(ChromaDB_VectorStore, ZhipuAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        ZhipuAI_Chat.__init__(self, config=config)


vn = MyVanna(

    config={"model": "glm-4", "api_key": "d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"})  # chatglm2-6b

vn.connect_to_mysql(host='10.240.53.162',
                    dbname='test',
                    user='rdsroot',
                    password='Geely@db20211206',
                    port=3306)
vn.train(ddl="""
CREATE TABLE IF NOT EXISTS users (
    id INT PRIMARY KEY COMMENT '用户ID'， 
    username VARCHAR(50) COMMENT '用户名'， 
    email VARCHAR(100) COMMENT '电子邮件'， 
    age INT COMMENT '年龄'， 
    gender VARCHAR(10) COMMENT '性别（男/女）'， 
    city VARCHAR(50) COMMENT '城市'
) COMMENT='用户信息表' CHARACTER SET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
""")

VannaFlaskApp(
    vn=vn,
    auth=SimplePassword(users=[{"email": "shen.5109256@qq.com", "password": "admin1"}]),
    allow_llm_to_see_data=True,
    title="Welcome to GouDan.AI",
    subtitle="Your AI-powered copilot for SQL queries.",
    show_training_data=True,
    sql=True,
    table=True,
    chart=True,
    summarization=False,
    ask_results_correct=True,
    # logo="https://account.geely.com//idaas/66526dd3c8fe6a8166886bcb7b3e39f6AA4TVXqcKoq"
    logo="https://i.geely.com/favicon.ico"
).run()
